import os
from torch.utils.data import DataSet
from transformers import Bert
from pathlib import Path
from typing import List
from collections import namedtuple
# 你要的类别集合
cls_ = ["类别a","类别b"]
# 先定义两个 类别映射后面处理数据集和转换分类结果用得到
name2id = {x:idx for idx,x in enumerate(cls_)}
id2name = {idx:x for idx,x in enumerate(cls_)}

data = namedtuple("data",["txt","label"])
class YourDataSet(DataSet):
    def __init__(self,dir:str,c):
        # 定义数据库文件夹目录
        p = Path(dir)
        if p.exists() and p.is_dir():
            self.filelist = [sub_p for sub_p in p.iterdir() if sub_p.is_file()]
            self.cls = c
            # 数据缓存列表
            self.db = []
            self.processData(self.filelist)
            self.lenght =len(self.db)
        else:
            raise ValueError(f"{dir} not found")
        
    # 这个函数不是必须重写的这里用来分离数据加载逻辑
    def processData(self,flist:List[Path]):
        for f in flist:
            # 这里检查文件名是不是在你的类别集合里面
            if f.stem() in self.cls:
                with open(f.absolute(),"r",encoding="utf-8") as wf:
                    # 按照行来取文本
                    for line in wf:
                        self.db.append(data(txt=f"{line.strip()}",label=name2id[f.stem()]))
    # 必须重写
    def __getitem__(self,idx):
        return self.db[idx]
    def __len__(self):
        return self.lenght
    
# 上面这部分就定义好了数据集了，如果训练集，评估集，测试集没有分多个文件夹，就得预处理拆分一下，下面主要是定义 `collate_fn`
def bertCollate(dl:List[data]):
    # 这里要返回一个batch(val,label) 的元组
    for d in dl:







